package com.sc.nft.helper;

import cn.hutool.core.date.DatePattern;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.StrUtil;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationContext;
import com.fasterxml.jackson.databind.JsonDeserializer;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.module.SimpleModule;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.lucene.search.TotalHits;
import org.elasticsearch.action.DocWriteRequest;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.get.MultiGetRequest;
import org.elasticsearch.action.get.MultiGetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.aggregations.Aggregation;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.metrics.ParsedSingleValueNumericMetricsAggregation;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.lang.NonNull;
import org.springframework.util.CollectionUtils;

import java.io.IOException;
import java.io.Serializable;
import java.math.BigDecimal;
import java.math.RoundingMode;
import java.text.SimpleDateFormat;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.time.format.DateTimeFormatter;
import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * Elasticsearch 工具类，增强 RestHighLevelClient 类方法（日志，其方法使用 final 修饰导致无法使用AOP增强，这里手动封装之）
 *
 * @author zlikun
 * @see RestHighLevelClient
 * @since 2024/06/06
 */
@Slf4j
public class ElasticsearchHelper {

    @Autowired
    private RestHighLevelClient client;

    private ObjectMapper mapper;

    /**
     * @see RestHighLevelClient#get(GetRequest, RequestOptions)
     */
    public GetResponse get(GetRequest request) throws IOException {
//        if (log.isDebugEnabled()) {
//            log.debug("GetRequest -> \nGET {}/_doc/{}", ArrayUtil.firstNonNull(request.indices()), request.id());
//        }
        return client.get(request, RequestOptions.DEFAULT);
    }

    public <T> T get(GetRequest request, Class<T> type) throws IOException {
        GetResponse response = get(request);
        if (!response.isExists()) {
            return null;
        }
        return mapper.readValue(response.getSourceAsString(), type);
    }

    public <T> T get(String index, String id, Class<T> type) throws IOException {
        return get(new GetRequest(index).id(id), type);
    }

    /**
     * @see RestHighLevelClient#search(SearchRequest, RequestOptions)
     */
    public SearchResponse search(SearchRequest request) throws IOException {
//        if (log.isDebugEnabled()) {
//            String query = request.source().toString();
//            if (query.contains("query")) {
//                log.debug("SearchRequest -> \nGET {}/_search\n{}", ArrayUtil.firstNonNull(request.indices()), query);
//            } else {
//                log.debug("SearchRequest -> \nGET {}/_search\n{\"query\": {}}", ArrayUtil.firstNonNull(request.indices()), query);
//            }
//        }
        return client.search(request, RequestOptions.DEFAULT);
    }

    /**
     * @see #search(SearchRequest)
     * @see RestHighLevelClient#search(SearchRequest, RequestOptions)
     */
    public <T> List<T> search(SearchRequest request, Class<T> type) throws IOException {
        SearchResponse response = search(request);
        return parseItems(response.getHits(), type);
    }

    /**
     * 搜索全部符合条件数据，规避ES查询上限（10000）
     *
     * @param index   查询索引
     * @param builder 查询条件
     * @param type    实体类型
     * @param <T>
     * @return
     * @throws IOException
     */
    public <T> List<T> searchAfter(String index, SearchSourceBuilder builder, Class<T> type) throws IOException {

        final int BATCH_SIZE = 1000; // 单页查询1000条
        long searchAfterValue = 0L; // 第一页，从0开始查询

        final List<T> container = new ArrayList<>();

        builder = builder.sort("id").size(BATCH_SIZE);
        while (true) {
            SearchSourceBuilder source = builder.searchAfter(new Object[]{searchAfterValue});
//            if (log.isDebugEnabled()) {
//                log.debug("SearchAfterRequest -> \nGET {}/_search\n{}", index, source);
//            }

            SearchRequest request = new SearchRequest(index).source(source);
            SearchResponse response = client.search(request, RequestOptions.DEFAULT);
            List<T> part = parseItems(response.getHits(), type);

            if (CollectionUtils.isEmpty(part)) {
                break;
            }

            container.addAll(part);

            if (part.size() < BATCH_SIZE) {
                break;
            }

            // 移动游标
            SearchHit[] hits = response.getHits().getHits();
            searchAfterValue = ((Number) hits[hits.length - 1].getSortValues()[0]).longValue();
        }

        return container;

    }

    /**
     * @see RestHighLevelClient#index(IndexRequest, RequestOptions)
     */
    public <T> IndexResponse index(@NonNull String index, @NonNull Serializable id, @NonNull T source) throws IOException {
        IndexRequest request = new IndexRequest(index);
        request.opType(DocWriteRequest.OpType.INDEX);
        request.id(Objects.requireNonNull(id).toString());
        request.source(ofObjectMapper().writeValueAsString(source), XContentType.JSON);
        return index(request);
    }

    /**
     * @see RestHighLevelClient#index(IndexRequest, RequestOptions)
     */
    private IndexResponse index(IndexRequest request) throws IOException {
//        if (log.isDebugEnabled()) {
//            log.debug("IndexRequest -> \nPUT {}/_doc/{}\n{}",
//                    ArrayUtil.firstNonNull(request.indices()), request.id(), request.source().utf8ToString());
//        }
        return client.index(request, RequestOptions.DEFAULT);
    }

    public <T> UpdateResponse update(@NonNull String index, @NonNull Serializable id, @NonNull T source) throws IOException {
        UpdateRequest request = new UpdateRequest(index, id.toString());
        request.doc(ofObjectMapper().writeValueAsString(source), XContentType.JSON);
        return update(request);
    }

    /**
     * @see RestHighLevelClient#update(UpdateRequest, RequestOptions)
     */
    private UpdateResponse update(UpdateRequest request) throws IOException {
//        if (log.isDebugEnabled()) {
//            log.debug("UpdateRequest -> \nPOST {}/_update/{}\n{ \"doc\": {} }",
//                    ArrayUtil.firstNonNull(request.indices()), request.id(), request.doc().source().utf8ToString());
//        }
        return client.update(request, RequestOptions.DEFAULT);
    }

    public DeleteResponse delete(@NonNull String index, @NonNull Serializable id) throws IOException {
        return delete(new DeleteRequest(index, id.toString()));
    }

    /**
     * @see RestHighLevelClient#delete(DeleteRequest, RequestOptions)
     */
    private DeleteResponse delete(DeleteRequest request) throws IOException {
//        if (log.isDebugEnabled()) {
//            log.debug("DeleteRequest -> \nDELETE {}/_doc/{}", ArrayUtil.firstNonNull(request.indices()), request.id());
//        }
        return client.delete(request, RequestOptions.DEFAULT);
    }

    /**
     * @see RestHighLevelClient#count(CountRequest, RequestOptions)
     */
    public CountResponse count(CountRequest request) throws IOException {
//        if (log.isDebugEnabled()) {
//            log.debug("CountRequest -> \nGET {}/_count\n{\"query\": {}}", ArrayUtil.firstNonNull(request.indices()), request.query());
//        }
        return client.count(request, RequestOptions.DEFAULT);
    }

    /**
     * 提取记录总数
     */
    public Long parseTotal(SearchHits hits) {
        if (hits == null) {
            return null;
        }

        TotalHits total = hits.getTotalHits();
        if (total == null) {
            return null;
        }

        return total.value;
    }

    public <T> List<T> parseItems(SearchHits hits, Class<T> type) {
        return parseItems(hits, hit -> {
            try {
                return mapper.readValue(hit.getSourceAsString(), type);
            } catch (JsonProcessingException e) {
                throw new RuntimeException(e);
            }
        });
    }

    public <T> List<T> parseItems(SearchHits hits, Function<SearchHit, T> mapping) {
        if (hits == null) {
            return new ArrayList<>();
        }

        SearchHit[] items = hits.getHits();
        if (ArrayUtil.isEmpty(items)) {
            return new ArrayList<>();
        }

        return Arrays.stream(items).map(mapping).collect(Collectors.toList());
    }

    private final String dateFormatPattern = DatePattern.NORM_DATETIME_PATTERN;
    private final DateTimeFormatter formatter = DateTimeFormatter.ofPattern(dateFormatPattern);

    public String ofDateTimeFormatPattern() {
        return dateFormatPattern;
    }

    /**
     * 统一ES时间格式：yyyy-MM-dd HH:mm:ss
     */
    public DateTimeFormatter ofDateTimeFormatter() {
        return formatter;
    }

    /**
     * 格式化输出时间，格式：yyyy-MM-dd HH:mm:ss
     */
    public String ofDateTimeString(LocalDateTime localDateTime) {
        return ofDateTimeFormatter().format(localDateTime);
    }

    /**
     * 格式化输出时间，格式：yyyy-MM-dd HH:mm:ss
     */
    public String ofDateTimeString(LocalDate date, LocalTime time) {
        return ofDateTimeString(LocalDateTime.of(date, time));
    }

    /**
     * 格式化输出时间，格式：yyyy-MM-dd 00:00:00
     */
    public String ofDateTimeString(LocalDate date) {
        return ofDateTimeString(date, LocalTime.MIN);
    }

    public ObjectMapper ofObjectMapper() {
        return mapper;
    }

    @Autowired
    public void setMapper(ObjectMapper mapper) {
        this.mapper = mapper.copy();
        this.mapper.setDateFormat(new SimpleDateFormat(dateFormatPattern));
        // 序列化时忽略NULL值，避免更新时强制将NULL更新到ES中
        this.mapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
        // BigDecimal 固定精度值为2，解决ES浮点值部分0时只保留1位精度问题（实际造成与MySQL的值精度不匹配）
        this.mapper.registerModule(new SimpleModule().addDeserializer(BigDecimal.class, new JsonDeserializer<BigDecimal>() {
            @Override
            public BigDecimal deserialize(JsonParser parser, DeserializationContext context) throws IOException, JsonProcessingException {
                String val = parser.getText();
                BigDecimal decimal = new BigDecimal(val);
                return decimal.setScale(2, RoundingMode.HALF_UP);
            }
        }));
    }

    public <T> List<T> mget(String index, Collection<String> ids, Class<T> type) throws IOException {
        if (CollectionUtils.isEmpty(ids)) {
            return new ArrayList<>();
        }
//        if (log.isDebugEnabled()) {
//            log.debug("GET {}/_mget\n{ \"ids\": {} }", index, ids);
//        }
        MultiGetRequest request = new MultiGetRequest();
        ids.forEach(id -> request.add(index, id));
        MultiGetResponse response = client.mget(request, RequestOptions.DEFAULT);
        final List<T> list = new ArrayList<>();
        response.spliterator().forEachRemaining(item -> {
            if (item.getResponse().isExists()) {
                try {
                    list.add(mapper.readValue(item.getResponse().getSourceAsString(), type));
                } catch (Exception e) {
                    log.error("解析MGET响应出错 index: {}, ids: {}", index, ids, e);
                }
            }
        });
        return list;
    }

    public <T> BulkResponse bulkIndex(String index, Collection<T> items, Function<T, String> docIdGetter) throws IOException {
        final BulkRequest request = new BulkRequest(index);
        items.stream()
                .filter(item -> item != null && docIdGetter.apply(item) != null)
                .forEach(item -> request.add(
                                new IndexRequest()
                                        .id(docIdGetter.apply(item))
                                        .source(writeValueAsString(item), XContentType.JSON)
                        )
                );

        return client.bulk(request, RequestOptions.DEFAULT);
    }

    @SneakyThrows
    private String writeValueAsString(Object data) {
        return mapper.writeValueAsString(data);
    }

    /**
     * 聚合查询
     */
    public Aggregations aggregation(String index, QueryBuilder query, AggregationBuilder aggregation) throws IOException {

        SearchRequest request = new SearchRequest()
                .indices(index)
                .source(new SearchSourceBuilder()
                        .query(query)
                        .aggregation(aggregation)
                        .size(0)
                );

//        String source = request.source().toString();
//        if (source.contains("query")) {
//            log.debug("SearchRequest -> \nGET {}/_search\n{}", ArrayUtil.firstNonNull(request.indices()), source);
//        } else {
//            log.debug("SearchRequest -> \nGET {}/_search\n{\"query\": {}}", ArrayUtil.firstNonNull(request.indices()), source);
//        }

        SearchResponse response = client.search(request, RequestOptions.DEFAULT);

        return response.getAggregations();
    }

    public BigDecimal parseAggregationToBigDecimalValue(Aggregations aggs, String field) {
        Aggregation agg = aggs.get(field);
        if (agg == null) {
            return null;
        }
        String value = ((ParsedSingleValueNumericMetricsAggregation) agg).getValueAsString();
        if (StrUtil.isBlank(value) || StrUtil.equals(value, "-Infinity")) {
            return null;
        }
        return new BigDecimal(value).setScale(2, RoundingMode.HALF_UP);
    }

}
